#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <mpi.h>

#include "protocol.h"

typedef struct node {
	int source;
	int destination;
	int weight;
} Node;

int   **nodes;
int   **weights;
int   **external;
int   *arraySize;
Node  **connections;
int   totalConnections = 0;

void printNodes(int *node) {
	int i = 0;
	for(i = 1; i < node[0]; i = i+2) {
		printf("%d|%c:%d ", i, node[i], node[i+1]);
	}
}

void printWeights(int p, int *weight) {
	int i = 0;
	for(i = 1; i < weight[0]; i += 2) {
		if(weight[i] == -1)
			printf("%d|%d:%d ", i, weight[i], weight[i+1]);
		else
			printf("%d|%c:%d ", i, nodes[p][weight[i]], weight[i+1]);
	}
}

void printExternal(int *extrn) {
	int i = 0;
	for(i = 1; i < extrn[0]; i += 3) {
		printf("%d|%c:%d:%d ", i, extrn[i], extrn[i+1], extrn[i+2]);
	}
}

void readNodes(FILE *inputFile, int *node) {
	int i = 0;
	for(i = 1; i < node[0]; i = i+2) {
		fscanf(inputFile, "%c\n", node+i);
	}
}

void readConnections(FILE *inputFile) {
	int src = 0;
	int dest = 0;
	int weight = 0 ;
	int count = 0;
	fscanf(inputFile, "%d\n", &count);

	// These are bi-directional graphs
	connections = malloc(count*sizeof(Node*)*2);
	totalConnections = count*2;

	int i = 0;
	for(i = 0; i+1 < totalConnections; i = i+2) {
		fscanf(inputFile, "%c %c %d\n", &src, &dest, &weight);
		connections[i] = malloc(sizeof(Node));
		connections[i]->source = src;
		connections[i]->destination = dest;
		connections[i]->weight = weight;
		connections[i+1] = malloc(sizeof(Node));
		connections[i+1]->source = dest;
		connections[i+1]->destination = src;
		connections[i+1]->weight = weight;
	}
}

// Finds the process that owns a given node
int findOwner(int procCount, int dest) {
	int i = 0;
	for(i = 0; i < procCount; i++) {
		int j = 0;
		for(j = 1; j < nodes[i][0]; j += 2) {
			if(nodes[i][j] == dest) {
				return i+1;
			}
		}
	}
	return -1;
}

void findConnections(int procCount) {
	int weightBuff[totalConnections*2];
	int externBuff[totalConnections*3];
	int weightCount = 0;
	int externCount = 0;

	// Loop for each process
	int proc = 0;
	for(proc = 0; proc < procCount; proc++) {
		// Slot 0 is reserved for length
		weightCount = 1;
		externCount = 1;
		
		// Loop over each node owned by the process
		int i = 0;
		for(i = 1; i < nodes[proc][0]; i = i+2) {
			// Store the index of first weight
			nodes[proc][i+1] = weightCount;
			// Find all connections of that node
			int j = 0;
			for(j = 0; j < totalConnections; j++) { 
				// If the connection source is owened by process
				if(connections[j]->source == nodes[proc][i]) {
					// Check if destination is owned
					_Bool owned = 0;
					int k = 0;
					for(k = 1; k < nodes[proc][0]; k = k+2) {
						if(connections[j]->destination == nodes[proc][k]) {
							owned = 1;
							break;
						}
					}
					if(owned) {
						weightBuff[weightCount] = k;
						weightBuff[weightCount+1] = connections[j]->weight;
						weightCount += 2;
					}else {
						weightBuff[weightCount] = -1;
						weightBuff[weightCount+1] = externCount;
						weightCount += 2;

						// Find the proc that owns the node
						int owner = findOwner(procCount, connections[j]->destination);
						if(owner == -1) {
							fprintf(stderr, "No owner for node %c found\n", 
									  connections[j]->destination);
							exit(13);
						}

						externBuff[externCount] = connections[j]->destination;
						externBuff[externCount+1] = owner;
						externBuff[externCount+2] = connections[j]->weight;
						externCount += 3;
					}
				}
			}
		}

		// Store length in array
		weightBuff[0] = weightCount;
		externBuff[0] = externCount;

		// Copy buffer into properly sized array
		weights[proc] = malloc(weightCount*sizeof(int));
		external[proc] = malloc(externCount*sizeof(int));

		memcpy(weights[proc], weightBuff, weightCount*sizeof(int));
		memcpy(external[proc], externBuff, externCount*sizeof(int));
	}
}

int main(int argc, char** argv) {
	char  *fileName = "test1.graph";
	FILE  *inputFile;
	int procCount = 6;
	int nodeCount = 0;
	int nodePerProc = 0;
	int remainder = 0;
	int Rank, Size;
	int scratch;
	_Bool *procDone;
	int flag;
	MPI_Status Status;

	MPI_Init ( &argc, &argv );
	MPI_Comm_rank ( MPI_COMM_WORLD, &Rank );
	MPI_Comm_size ( MPI_COMM_WORLD, &Size );


	procCount = Size-1;
	
	if(Rank == 0)
	{
		fileName = argc > 1 ? argv[1] : fileName;
		inputFile = fopen(fileName, "r");

		if ( inputFile == NULL )
		{  fprintf (stderr, "Open failed for %s\n", fileName); exit(666);  }
		
		fscanf(inputFile, "%d\n", &nodeCount);

		nodePerProc = nodeCount/procCount;
		remainder = nodeCount%procCount;

		/*
		 * Creat the NODES array for all processors
		 */
		nodes = malloc(procCount*sizeof(int*));
		int i = 0;
		for(i = 0; i < procCount; i++) {
			int count = nodePerProc + ((remainder-- > 0) ? 1 : 0);
			// Add one to hold the length of the array
			nodes[i] = calloc((1+count*2),sizeof(int));
			nodes[i][0] = 1+count*2;
			readNodes(inputFile, nodes[i]);
		}

		readConnections(inputFile);
		fclose(inputFile);

		// Construct the weights and external array
		weights = malloc(procCount*sizeof(int*));
		external = malloc(procCount*sizeof(int*));

		findConnections(procCount);

		/*
		 * Print the node array
		 * Print the weight array
		 * Print the external array
		 */
		for(i = 0; i < procCount; i++) {
			printf("Process %d\n", i+1);
			printNodes(nodes[i]);
			printf("\n");
			printWeights(i, weights[i]);
			printf("\n");
			printExternal(external[i]);
			printf("\n\n");
		}

		arraySize = malloc(3*sizeof(int));
		for(i = 0; i < procCount; i++)
		{
			arraySize[0] = nodes[i][0];
			arraySize[1] = weights[i][0];
			arraySize[2] = external[i][0];
			MPI_Send(arraySize, 3, MPI_INT, i+1, INIT, MPI_COMM_WORLD);
			MPI_Send(nodes[i], arraySize[0], MPI_INT, i+1, NODES, MPI_COMM_WORLD);
			MPI_Send(weights[i], arraySize[1], MPI_INT, i+1, WEIGHTS, MPI_COMM_WORLD);
			MPI_Send(external[i], arraySize[2], MPI_INT, i+1, EXTERNAL, MPI_COMM_WORLD);
		}

		// Send the first Worker the starting node
		int nw[3];
		nw[0] = nodes[0][1];
		nw[1] = 0;
		nw[2] = ' ';
		MPI_Send(nw, 3, MPI_INT, 1, NW, MPI_COMM_WORLD);

		procDone = malloc(procCount*sizeof(_Bool));

		for(i = 0 ; i < procCount ; i++)
		{
			procDone[i] = 0;
		}
		_Bool working = 1;
		_Bool notDone = 1;
		_Bool allDone = 1;
		_Bool checkForUndone = 1;
		/*
		MPI_Status is defined:
			int MPI_SOURCE;
			int MPI_TAG;
			int MPI_ERROR;
		} MPI_Status;
		*/
		while(working)
		{
			while(notDone)
			{
				allDone = 1;
if(DEBUG) printf("0 Waiting on DONE\n");
				MPI_Recv(&scratch, 0, MPI_INT, MPI_ANY_SOURCE, DONE, MPI_COMM_WORLD, &Status);
				procDone[Status.MPI_SOURCE-1] = 1;
if(DEBUG) printf("0 Done from: %d\n", Status.MPI_SOURCE);
				for(i = 0; i < procCount ; i++)
				{
					if(!procDone[i])
						allDone = 0;
				}
				if(allDone) {
if(DEBUG) printf("0 All done\n");
					notDone = 0;
				}
			}
			for(i = 0 ; i < procCount ; i++)
			{
				MPI_Send(&scratch, 0, MPI_INT, i+1, FINALIZE, MPI_COMM_WORLD);
			}
if(DEBUG) printf("0 Finalize Sent\n");
			allDone = 1;
if(DEBUG) printf("0 Before UNDONE Loop\n");
			for(i = 0 ; i < procCount ; i++)
			{
//if(DEBUG) printf("In check UNDONE loop\n");
				MPI_Recv(&scratch, 0, MPI_INT, i+1, MPI_ANY_TAG, MPI_COMM_WORLD, &Status);
if(DEBUG) printf("0 Recieved from: %d\n", i+1);
				if(Status.MPI_TAG == UNDONE)
				{
if(DEBUG) printf("0 Recieved UNDONE from: %d\n", i+1);
					procDone[i] = 0;
				}
				else if(Status.MPI_TAG == DONE)
				{
if(DEBUG) printf("0 Recieved DONE from: %d\n", i+1);
					procDone[i] = 1;
				}
			}
			for(i = 0; i < procCount ; i++)
			{
				if(!procDone[i]){
if(DEBUG) printf("\033[22;32m0 proc not done: %d\n\033[22;37m", i+1);
					allDone = 0;
                                }
			}
			if(allDone) {
if(DEBUG) printf("0 All done After Finalize\n");
				working = 0;
			}
		}
		for(i = 0 ; i < procCount ; i++)
		{
			MPI_Send(&scratch, 0, MPI_INT, i+1, REPORT, MPI_COMM_WORLD);
		}
if(DEBUG) printf("0 Report Sent\n");
		sleep(2);
		int j;
		for(i = 0 ; i < procCount ; i++)
		{
if(DEBUG) printf("0 In Report Loop\n");
			int somesize = 0;
			MPI_Recv(&somesize, 1, MPI_INT, i+1, REPORTSIZE, MPI_COMM_WORLD, &Status);
if(DEBUG) printf("0 Recieving ReportSize from %d\n", i+1);
			int *someArray;
			someArray = malloc(somesize*sizeof(int));
			MPI_Recv(someArray, somesize, MPI_INT, i+1, REPORT, MPI_COMM_WORLD, &Status);
		// Conglomerate info
if(DEBUG) printf("0 Recieved Report from %d %d\n", i+1, somesize);
			for(j = 0; j < somesize; j += 3) {
				printf("Destination: %c\n", someArray[j]);
				printf("Distance   : %d\n", someArray[j+1]);
				printf("Previous   : %c\n", someArray[j+2]);
			}
			free(someArray);
		}	

		MPI_Barrier(MPI_COMM_WORLD);
		// Free memory
		for(i = 0; i < procCount; i++) {
			free(nodes[i]);
			free(weights[i]);
			free(external[i]);
		}
		free(procDone);
		free(arraySize);
		free(nodes);
		free(weights);
		free(external);

		for(i = 0; i < totalConnections; i++) {
			free(connections[i]);
		}
		free(connections);
	}
	MPI_Finalize();
	return 0;
}
